# 针对一个路径下的一组dicom文件。 包括了stack、mpr。
import os
from dataclasses import dataclass

import time
from components.Viewer import Viewer
from components.CoordTransfer import CoordTransfer
from components.VolumeInfo import VolumeInfo
from components.dicom_tag_info import DicomTagInfo
from connect import send_message
from type import EnumViewType, EnumOpType, EnumImgQuality, EnumImgMode
from imgLoader.VolumeLoader import VolumeLoader
import asyncio
from utils.DicomLoader import DicomLoader
from utils.util import debounce
from share import watching_dicom_queue_dict

load_viewer_arr = [EnumViewType.Ax.value, EnumViewType.Sag.value, EnumViewType.Cor.value]


@dataclass
class LoaderState:
    has_dicom = False
    stack_size_set = False
    stack_transform_init = False

    def set_has_dicom_state(self, state):
        self.has_dicom = state

    def set_stack_size_set_state(self, state):
        self.stack_size_set = state

    def set_stack_transform_init_state(self, state):
        self.stack_transform_init = state


class ImgLoader:
    def __init__(self, client_id, dicom_folder_path):
        self.client_id = client_id
        self.dicom_folder_path = dicom_folder_path
        self.dicom_tag_info = DicomTagInfo()
        self.coord_transfer: CoordTransfer = CoordTransfer()
        self.volume_info = VolumeInfo(self.dicom_tag_info, self.coord_transfer)
        self.stack_viewer = Viewer(view_type=EnumViewType.Stack.value, dicom_tag_info=self.dicom_tag_info,
                                   volume_info=self.volume_info)
        self.ax_viewer = Viewer(view_type=EnumViewType.Ax.value, dicom_tag_info=self.dicom_tag_info,
                                volume_info=self.volume_info)
        self.sag_viewer = Viewer(view_type=EnumViewType.Sag.value, dicom_tag_info=self.dicom_tag_info,
                                 volume_info=self.volume_info)
        self.cor_viewer = Viewer(view_type=EnumViewType.Cor.value, dicom_tag_info=self.dicom_tag_info,
                                 volume_info=self.volume_info)
        self.viewers = []
        self.volume_viewer = VolumeLoader(self.dicom_tag_info)
        self.phase_index_to_ds_arr_dict = {}
        self.loader_state = LoaderState()
        self.is_playing = False
        self.dicom_loader = DicomLoader(dicom_folder_path)

        if load_viewer_arr.__contains__(EnumViewType.Sag.value):
            self.viewers.append({
                "view_type": EnumViewType.Sag.value,
                "viewer": self.sag_viewer
            })
        if load_viewer_arr.__contains__(EnumViewType.Cor.value):
            self.viewers.append({
                "view_type": EnumViewType.Cor.value,
                "viewer": self.cor_viewer
            })
        if load_viewer_arr.__contains__(EnumViewType.Ax.value):
            self.viewers.append({
                "view_type": EnumViewType.Ax.value,
                "viewer": self.ax_viewer
            })

    async def load_dicom_at_once(self):
        print('读取dicom文件夹')
        self.phase_index_to_ds_arr_dict, first_phase_index = self.dicom_loader.load_a_exist_fold(self.dicom_folder_path)
        print(f'dicom文件夹读取完成，长度为：{len(self.phase_index_to_ds_arr_dict)}')
        if len(self.phase_index_to_ds_arr_dict) == 0:
            return
        self.dicom_tag_info.init_tag_info(self.phase_index_to_ds_arr_dict[first_phase_index][0])
        self.volume_info.init_spacing_info()
        self.init_after_volume_loaded()
        self.loader_state.set_has_dicom_state(True)
        await self.send_current_stack_img()

    def init_after_volume_loaded(self):
        self.volume_info.init_origin_volume_by_exist_files(self.phase_index_to_ds_arr_dict)
        self.volume_info.init_stack_view_current_index()

    def init_mpr(self):
        self.volume_info.init_norm_volume()
        self.volume_info.init_mpr_view_current_index()

    async def start_watch_queue(self):
        q = watching_dicom_queue_dict.get(self.client_id)
        while True:
            if q.qsize() > 0:
                dicom_path = q.get()

                await self.load_a_dicom(dicom_path)

            else:
                await asyncio.sleep(1)

    async def load_a_dicom(self, dicom_file_path):
        new_slice, phase_index = self.dicom_loader.load_a_dicom(dicom_file_path)
        if phase_index in self.phase_index_to_ds_arr_dict:
            self.phase_index_to_ds_arr_dict[phase_index].append(new_slice)
        else:
            self.phase_index_to_ds_arr_dict[phase_index] = [new_slice]

        if self.loader_state.has_dicom is False:
            self.dicom_tag_info.init_tag_info(new_slice)
            self.volume_info.init_spacing_info()
            pixel_array = [new_slice.pixel_array]
            self.volume_info.init_origin_volume_by_coming_file(pixel_array, phase_index)
            print('init volume set finish')
        else:
            print('add a slice')
            self.volume_info.add_a_slice(new_slice.pixel_array, phase_index)
        self.loader_state.set_has_dicom_state(True)
        await self.send_current_stack_img()

    async def send_current_stack_img(self):
        if self.loader_state.stack_size_set is False or self.loader_state.has_dicom is False:
            return
        if self.loader_state.stack_transform_init is False:
            self.volume_info.init_transform([EnumViewType.Stack.value])
            self.loader_state.set_stack_transform_init_state(True)
        img_dict = {}
        img_info = self.stack_viewer.get_recent_img(EnumImgQuality.low.value)
        img_dict.setdefault(EnumViewType.Stack.value, img_info)
        await send_message(self.client_id, img_dict)
        await self.delayed_render(EnumViewType.Stack.value, '')

    async def set_stack_size(self, size):
        print('set stack size')
        self.volume_info.set_canvas_size(EnumViewType.Stack.value, size)
        self.loader_state.set_stack_size_set_state(True)
        await self.send_current_stack_img()

    async def set_mpr_viewer_size(self, size):
        first_phase_index = list(self.phase_index_to_ds_arr_dict.keys())[0]
        self.coord_transfer.init(self.phase_index_to_ds_arr_dict[first_phase_index])
        self.init_mpr()

        ax_size = (size[EnumViewType.Ax.value]['height'], size[EnumViewType.Ax.value]['width'])
        sag_size = (size[EnumViewType.Sag.value]['height'], size[EnumViewType.Sag.value]['width'])
        cor_size = (size[EnumViewType.Cor.value]['height'], size[EnumViewType.Cor.value]['width'])
        self.volume_info.set_canvas_size(EnumViewType.Ax.value, ax_size)
        self.volume_info.set_canvas_size(EnumViewType.Sag.value, sag_size)
        self.volume_info.set_canvas_size(EnumViewType.Cor.value, cor_size)
        self.volume_info.init_transform([
            EnumViewType.Ax.value,
            EnumViewType.Sag.value,
            EnumViewType.Cor.value
        ])
        await self.get_view(EnumViewType.MPR.value)

    def set_volume_viewer_size(self, size):
        self.volume_viewer.load_dicom_series(self.dicom_folder_path)
        self.volume_viewer.set_render_size(size)

    async def get_view(self, view_type: EnumViewType):
        ans = {}
        match view_type:
            case EnumViewType.Stack.value:
                img_info = self.stack_viewer.get_mix_img(EnumImgQuality.high.value)
                ans = {
                    EnumViewType.Stack.value: img_info
                }
            case EnumViewType.MPR.value:
                for each_viewer in self.viewers:
                    viewer = each_viewer['viewer']
                    view_type = each_viewer['view_type']
                    img_info = viewer.get_mix_img(EnumImgQuality.high.value)
                    ans.setdefault(view_type, img_info)
            case EnumViewType.Volume.value:
                img_info = self.volume_viewer.get_high_quality_img()
                ans = {
                    EnumViewType.Volume.value: img_info
                }
        await send_message(self.client_id, ans)

    def get_view_type(self, view_type):
        Stack = {
            "view_type": EnumViewType.Stack.value,
            "viewer": self.stack_viewer
        }
        Ax = {
            "view_type": EnumViewType.Ax.value,
            "viewer": self.ax_viewer
        }
        Sag = {
            "view_type": EnumViewType.Sag.value,
            "viewer": self.sag_viewer
        }
        Cor = {
            "view_type": EnumViewType.Cor.value,
            "viewer": self.cor_viewer
        }
        MPR = [Ax, Sag, Cor]
        match view_type:
            case EnumViewType.Stack.value:
                return [Stack]
            case EnumViewType.Sag.value:
                return [Sag]
            case EnumViewType.Ax.value:
                return [Ax]
            case EnumViewType.Cor.value:
                return [Cor]
            case EnumViewType.MPR.value:
                return MPR
            case EnumViewType.Volume.value:
                return [{
                    "view_type": 'volume',
                    "viewer": self.volume_viewer
                }]

    async def handle_op(self, msg):
        view_type = msg['viewType']
        op_type = msg['opType']
        viewer_arr = self.get_view_type(view_type)
        img_dict = dict()
        match op_type:
            case EnumOpType.window.value:
                x = msg['x']
                y = msg['y']
                print(f'set window: {x}, {y}')
                for viewer_dict in viewer_arr:
                    viewer: Viewer = viewer_dict['viewer']
                    viewer_type = viewer_dict['view_type']
                    viewer.update_window(x, y)
                    img_info = viewer.get_recent_img(EnumImgQuality.low.value)
                    img_dict.setdefault(viewer_type, img_info)
            case EnumOpType.flip.value:
                hor_flip = msg['hor_flip']
                ver_flip = msg['ver_flip']
                for viewer_dict in viewer_arr:
                    viewer: Viewer = viewer_dict['viewer']
                    viewer_type = viewer_dict['view_type']
                    viewer.update_flip(hor_flip, ver_flip)
                    img_info = viewer.get_recent_img(EnumImgQuality.low.value)
                    img_dict.setdefault(viewer_type, img_info)
            case EnumOpType.rotate.value:
                rotate_type = msg['subOpType']
                for viewer_dict in viewer_arr:
                    viewer: Viewer = viewer_dict['viewer']
                    viewer_type = viewer_dict['view_type']
                    viewer.update_rotate(rotate_type)
                    img_info = viewer.get_recent_img()
                    img_dict.setdefault(viewer_type, img_info)
            case EnumOpType.rotate3D.value:
                x = msg['x']
                y = msg['y']
                for viewer_dict in viewer_arr:
                    viewer: Viewer = viewer_dict['viewer']
                    viewer_type = viewer_dict['view_type']
                    viewer: VolumeLoader
                    viewer.update_rotate(x, y)
                    img_info = viewer.get_recent_img()

                    img_dict.setdefault(viewer_type, img_info)
            case EnumOpType.resize.value:
                for viewer_dict in viewer_arr:
                    viewer: Viewer = viewer_dict['viewer']
                    each_view_type = viewer_dict['view_type']
                    width = msg['size'][each_view_type]['width']
                    height = msg['size'][each_view_type]['height']
                    size = (height, width)
                    viewer.update_resize(size)
                    img_info = viewer.get_recent_img(EnumImgQuality.low.value)
                    img_dict.setdefault(each_view_type, img_info)
            case EnumOpType.getPosition.value:
                x = msg['x']
                y = msg['y']
                for viewer_dict in viewer_arr:
                    viewer: Viewer = viewer_dict['viewer']
                    viewer_type = viewer_dict['view_type']
                    viewer.get_current_mouse_position(x, y)
                    img_info = viewer.get_recent_img(EnumImgQuality.low.value)
                    img_dict.setdefault(viewer_type, img_info)
            case EnumOpType.zoom.value:
                zoom_factor = msg['zoom']
                for viewer_dict in viewer_arr:
                    viewer = viewer_dict['viewer']
                    viewer_type = viewer_dict['view_type']
                    viewer.update_zoom(zoom_factor)
                    img_info = viewer.get_recent_img()
                    img_dict.setdefault(viewer_type, img_info)
            case EnumOpType.pan.value:
                delta_x = msg['x']
                delta_y = msg['y']
                if delta_x == 0 and delta_y == 0:
                    pass
                else:
                    for viewer_dict in viewer_arr:
                        viewer = viewer_dict['viewer']
                        viewer_type = viewer_dict['view_type']
                        viewer.update_pan(delta_x, delta_y)
                        img_info = viewer.get_recent_img()
                        img_dict.setdefault(viewer_type, img_info)
            case EnumOpType.scroll.value:
                scroll = msg['scroll']
                for viewer_dict in viewer_arr:
                    viewer = viewer_dict['viewer']
                    viewer_type = viewer_dict['view_type']
                    viewer.update_current_slice(scroll)
                    if view_type == EnumViewType.Stack.value:
                        img_info = viewer.get_recent_img(EnumImgQuality.low.value)
                        img_dict.setdefault(viewer_type, img_info)

                if view_type != EnumViewType.Stack.value:
                    all_viewer_arr = self.get_view_type(EnumViewType.MPR.value)
                    for viewer_dict in all_viewer_arr:
                        viewer = viewer_dict['viewer']
                        viewer_type = viewer_dict['view_type']
                        img_info = viewer.get_recent_img(EnumImgQuality.low.value)
                        img_dict.setdefault(viewer_type, img_info)
            case EnumOpType.scrollTo.value:
                scroll_to_index = msg['toIndex']
                for viewer_dict in viewer_arr:
                    viewer = viewer_dict['viewer']
                    viewer_type = viewer_dict['view_type']
                    viewer.slice_to_index(scroll_to_index)
                    img_info = viewer.get_recent_img(EnumImgQuality.low.value)
                    img_dict.setdefault(viewer_type, img_info)
            # 移动mpr中的cross
            case EnumOpType.move_scroll_absolute.value:
                x = msg['x']
                y = msg['y']
                self.volume_info.update_slice_position(view_type, x, y)
                for viewer_dict in self.viewers:
                    viewer = viewer_dict['viewer']
                    viewer_type = viewer_dict['view_type']
                    img_info = viewer.get_recent_img(EnumImgQuality.low.value)
                    img_dict.setdefault(viewer_type, img_info)
            case EnumOpType.measure.value:
                sub_op_type = msg['subOpType']
                x = msg['x']
                y = msg['y']
                action_type = msg['actionType']
                img_dict = dict()
                for viewer_dict in viewer_arr:
                    viewer: Viewer = viewer_dict['viewer']
                    viewer_type = viewer_dict['view_type']
                    viewer.handle_draw(sub_op_type, (x, y), action_type)
                    img_info = viewer.get_recent_img(EnumImgQuality.low.value)
                    img_dict.setdefault(viewer_type, img_info)
            case EnumOpType.performance.value:
                sub_op_type = msg['subOpType']
                x = msg['x']
                y = msg['y']
                action_type = msg['actionType']
                img_dict = dict()
                for viewer_dict in viewer_arr:
                    viewer: Viewer = viewer_dict['viewer']
                    viewer_type = viewer_dict['view_type']
                    viewer.handle_draw(sub_op_type, (x, y), action_type)
                    img_info = viewer.get_recent_img(EnumImgQuality.low.value)
                    img_dict.setdefault(viewer_type, img_info)
            case EnumOpType.reset.value:
                reset_type = msg['subOpType']
                if reset_type == 'resetAll' and view_type == EnumViewType.MPR.value:
                    self.volume_info.reset_mpr_slice()
                for viewer_dict in viewer_arr:
                    viewer: Viewer = viewer_dict['viewer']
                    viewer_type = viewer_dict['view_type']
                    viewer.reset(reset_type)
                    img_info = viewer.get_recent_img()
                    img_dict.setdefault(viewer_type, img_info)

        await self.delayed_render(view_type, op_type)
        await send_message(self.client_id, img_dict)

    async def handle_mip(self, min_setting):
        useMIP = min_setting['useMIP']
        mip_option = min_setting['mipOption']
        mip_model = mip_option['model']
        mip_val = mip_option['value']
        self.volume_info.set_mip(useMIP, mip_model, mip_val)
        viewer_arr = self.get_view_type(EnumViewType.MPR.value)
        img_dict = dict()
        for viewer_dict in viewer_arr:
            viewer: Viewer = viewer_dict['viewer']
            viewer_type = viewer_dict['view_type']
            img_info = viewer.get_recent_img()
            img_dict.setdefault(viewer_type, img_info)
        await send_message(self.client_id, img_dict)

    async def handle_play(self, play_setting, is_4d=False):
        play_speed = play_setting['playSpeed']
        view_type = play_setting['viewType']
        self.is_playing = True
        if view_type == EnumViewType.Stack.value:
            viewer_arr = self.get_view_type(view_type)
        else:
            viewer_arr = self.get_view_type(EnumViewType.MPR.value)
        while True:
            if self.is_playing is False:
                break
            img_dict = {}
            self.volume_info.toggle_next(view_type, is_4d)
            for viewer_dict in viewer_arr:
                viewer: Viewer = viewer_dict['viewer']
                viewer_type = viewer_dict['view_type']
                img_info = viewer.get_recent_img()
                img_dict.setdefault(viewer_type, img_info)
            await send_message(self.client_id, img_dict)
            time.sleep(play_speed)

    def handle_pause(self):
        self.is_playing = False

    async def handle_phase_change(self, json_request):
        to_phase_index = json_request['toPhaseIndex']
        view_type = json_request['viewType']
        self.volume_info.phase_data.set_current(to_phase_index)
        match view_type:
            case EnumViewType.Stack.value:
                img_info = self.stack_viewer.get_recent_img(EnumImgQuality.high.value)
                img_dict = {
                    EnumViewType.Stack.value: img_info
                }
                await send_message(self.client_id, img_dict)

            case _:
                all_viewer_arr = self.get_view_type(EnumViewType.MPR.value)
                img_dict = {}
                for viewer_dict in all_viewer_arr:
                    viewer = viewer_dict['viewer']
                    viewer_type = viewer_dict['view_type']
                    img_info = viewer.get_recent_img(EnumImgQuality.high.value)
                    img_dict.setdefault(viewer_type, img_info)
                await send_message(self.client_id, img_dict)

    @debounce(0.3)
    async def delayed_render(self, view_type, op_type):
        viewer_arr = self.get_view_type(view_type)
        img_dict = {}
        if op_type == EnumOpType.scroll.value and view_type != EnumViewType.Stack.value and view_type != EnumViewType.Volume.value:
            all_viewer_arr = self.get_view_type(EnumViewType.MPR.value)
            for viewer_dict in all_viewer_arr:
                viewer = viewer_dict['viewer']
                viewer_type = viewer_dict['view_type']
                img_info = viewer.get_recent_img(EnumImgQuality.high.value)
                img_dict.setdefault(viewer_type, img_info)
        elif view_type == EnumViewType.Volume.value:
            img_info = self.volume_viewer.get_high_quality_img()
            img_dict.setdefault(view_type, img_info)
        else:
            for viewer_dict in viewer_arr:
                viewer: Viewer = viewer_dict['viewer']
                viewer_type = viewer_dict['view_type']
                img_info = viewer.get_recent_img(EnumImgQuality.high.value)
                img_dict.setdefault(viewer_type, img_info)
        print(f'发送延迟渲染图像，{view_type}')
        await send_message(self.client_id, img_dict)

    async def set_3d_params(self, json_request):
        template = json_request['template']
        windows = json_request['windows']
        self.volume_viewer.set_3d_params(template, windows)
        img_info = self.volume_viewer.get_high_quality_img()
        ans = {
            EnumViewType.Volume.value: img_info
        }
        return ans

    async def set_3d_direction(self, json_request):
        direction = json_request['direction']
        self.volume_viewer.set_3d_direction(direction)
        img_info = self.volume_viewer.get_high_quality_img()
        ans = {
            EnumViewType.Volume.value: img_info
        }
        return ans

    async def handle_performance(self, json_request):
        performance_mode = json_request["performance"]
        if performance_mode == '':
            self.stack_viewer.toggle_mode(EnumImgMode.img.value)
        elif performance_mode == 'mtf':
            self.stack_viewer.toggle_mode(EnumImgMode.performance.value)


